#coding=utf-8

# #  reg的自动定位；

import cv2
import numpy as np
import math
from colorCvt import  *
import os
import shutil

imgPath = "444.jpg"

SHOW_SAVE_IMG = False
SHOW_REG = False

# 可调餐宿；
ratio = 0.005
thresold = 0.9                                    # 初步检测的阈值；
roiSize = 40                                        # roi分析大小；

# 检测过程参数；
subThresold=[8000 , 8000, 1000, 1000]

def isReg(binaryImg, pos, bFlags = True):
    """
    判断是否是reg标记；         同时， 应该是黑白区域；
    :param binaryImg:
    :param pos:
    :return:
    """
    radius = [1.0 * roiSize/4.0, 1.0 * roiSize/3.0, roiSize/2.0 - 1]
    angle = [30, 45, 60]
    flags = []
    if bFlags:
        flags = [255, 0, 255, 0]
    else:
        flags = [0, 255, 0, 255]
    negPos = [[1, -1], [-1, -1], [-1, 1], [1, 1]]
    y = pos[0]
    x = pos[1]
    num = 0
    height = binaryImg.shape[0]
    width = binaryImg.shape[1]
    for k in range(4):                  # 四个象限；
        for r in range(len(radius)):
            tmpRadius = radius[r]
            for ag in range(len(angle)):
                tmpAngle = angle[ag]
                newX = int(1.0 * tmpRadius * math.cos(1.0 * tmpAngle/180 * math.pi)) * negPos[k][0] + x
                newY = int(1.0 * tmpRadius * math.sin(1.0 * tmpAngle/180 * math.pi)) * negPos[k][1] + y
                if newX >= width or newY >= height or newX < 0 or newY < 0:
                    continue
                if(binaryImg[newY][newX] == flags[k]):
                    num = num + 1
    nTotal = len(radius) * len(angle) * 4
    if 1.0 * num/nTotal > thresold:
        return True
    return False

def CreateImg(width, height, dim):
    """
    创建图像； 默认是255灰度级；
    :param width:
    :param height:
    :param dim:
    :return:
    """
    img = np.zeros([height, width, dim], np.uint8)
    return img

def SaveImg(img, path):
    """
    保存图像；
    :param img:
    :param path:
    :return:
    """
    cv2.imwrite(path, img)

def RGB2HSV(rgb):
    """
    RGB转为HSV；
    :param rgb:
    :return:
    """
    b, g, r = rgb[0]/255.0, rgb[1]/255.0, rgb[2]/255.0
    cmax = np.max((r, g, b))
    cmin = np.min((r, g, b))
    delta = cmax - cmin
    h = 0
    if cmax == r and g >= b:
        h = 60.0 * (g - b)/delta
    elif cmax == r and g < b:
        h = 60 * (g - b)/delta + 360
    elif cmax == g:
        h = 60.0 * ((b-r)/delta + 2)
    elif cmax == b:
        h = 60.0 * ((r - g)/delta + 4)
    else:
        h = 0

    l = 0.5 * (cmax + cmin)
    s = 0
    if l == 0 or cmax == cmin:
        s = 0
    elif l > 0 and l <= 0.5:
        s = (cmax -cmin)/(2.0 * l)
    elif l > 0.5:
        s = (cmax - cmin)/(2.0 * (1.0 - l))

    v = cmax
    return (h, s, v)

def ColorMap(img):
    """
    颜色谱；
    :param img:
    :return:
    """
    width = img.shape[1]
    height = img.shape[0]
    channel = img.shape[2]
    saveImg = CreateImg(width, height, 1)
    for  j in range(height):
        for i in range(width):
            pixel = img[j][i]
            r, g, b = pixel[2], pixel[1], pixel[0]
            x, y, z = rgb2xyz((r, g, b))
            l, a, b = xyz2lab((x, y, z))
            if a > 5:
                saveImg[j][i] = 255
    SaveImg(saveImg, "tt.jpg")

def AnalysisRegNear(contour, srcImg):
    """
    分析reg周边， 采用获取roi的方法， 再根据特征进行分析；
   :param contour:
    :param img:
    :return:
    """
    height = srcImg.shape[0]
    width = srcImg.shape[1]
    center = np.array([0] * 2, dtype="float64")              #  列；行;
    for i in range(len(contour)):
        y, x = contour[i][0]                # y为行， x为列；
        center = center + [x, y]
    center = center/(len(contour))
    # 边界处理；
    if center[1] - roiSize/2.0 < 0 or center[0] - roiSize/2.0 < 0:
        return  False
    if center[1] + roiSize/2.0 > width - 1 or center[0] + roiSize/2.0 > height - 1:
        return False

    roiImg = srcImg[int(center[0] - roiSize/2.0):int(center[0] + roiSize/2.0), int(center[1] - roiSize/2.0):int(center[1] + roiSize/2.0)]

    roiImg = cv2.cvtColor(roiImg, cv2.COLOR_BGR2GRAY)
    # roiImg =cv2.normalize(roiImg,dst=None,alpha=350,beta=10,norm_type=cv2.NORM_MINMAX)

    # 平均像素；
    kernel2 = cv2.getStructuringElement(cv2.MORPH_CROSS, (5, 5))
    roiImg = cv2.morphologyEx(roiImg, cv2.MORPH_CLOSE, kernel2)
    avgPixel = cv2.mean(roiImg)
    # print(avgPixel)

    # 阈值处理；；
    binaryImg = cv2.threshold(roiImg, int(avgPixel[0]), 255, cv2.THRESH_BINARY)
    binaryImg = np.array(binaryImg[1], dtype="uint8")

    contours2, hierarchy = cv2.findContours(binaryImg, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    if len(contours2) > 5:
        return False

    if isReg(binaryImg, (int(roiSize/2.0), int(roiSize/2.0)), True) or isReg(binaryImg, (int(roiSize/2.0), int(roiSize/2.0)), False):
        #相似度判断；
        if SHOW_REG:
            tmpROiImg = srcImg[int(center[0] - roiSize):int(center[0] + roiSize), int(center[1] - roiSize):int(center[1] + roiSize)]
            cv2.namedWindow("roiImg", 0);
            cv2.resizeWindow("roiImg", 640, 480)
            cv2.imshow("roiImg", tmpROiImg)
            cv2.waitKey(0)
        splitX = int(roiSize/2.0)
        splitY = int(roiSize/2.0)
        num1 = 0        # 两个风格的；
        num2 = 0
        for j in range(roiSize):
            for i in range(roiSize):
                offsetX = j - splitX
                offsetY = i - splitY
                if offsetX * offsetY < 0:
                    if binaryImg[j][i] == 255:
                        num1 = num1 + 1
                    if binaryImg[j][i] == 0:
                        num2 = num2 + 1
                else:
                    if binaryImg[j][i] == 255:
                        num2 = num2 + 1
                    if binaryImg[j][i] == 0:
                        num1 = num1 + 1
        # print(num1, num2)
        if 1.0 * num1/num2 <= 5 and 1.0 * num2/num1 <= 5:
            return False

        # 颜色判断；
        y = int(center[0])
        x = int(center[1])
        offset = 6
        tmpRoi_1 = srcImg[(y - offset * 2):(y -offset), (x - offset * 2): (x - offset)]         # 左上；
        tmpRoi_1_avg = np.array(cv2.mean(tmpRoi_1))

        tmpRoi_2 = srcImg[(y - offset * 2):(y - offset), (x + offset): (x + offset * 2)]        # 右上；
        tmpRoi_2_avg = np.array(cv2.mean(tmpRoi_2))

        tmpRoi_3 = srcImg[(y + offset): (y + offset * 2), (x - offset * 2): (x - offset)]       # 左下；
        tmpRoi_3_avg = np.array(cv2.mean(tmpRoi_3))

        tmpRoi_4 = srcImg[(y + offset): (y + offset * 2), (x + offset): (x + offset * 2)]       # 右下；
        tmpRoi_4_avg = np.array(cv2.mean(tmpRoi_4))

        sub12 = np.sum((tmpRoi_1_avg - tmpRoi_2_avg) * (tmpRoi_1_avg - tmpRoi_2_avg))
        sub14 = np.sum((tmpRoi_1_avg - tmpRoi_4_avg) * (tmpRoi_1_avg - tmpRoi_4_avg))
        sub23 = np.sum((tmpRoi_2_avg - tmpRoi_3_avg) * (tmpRoi_2_avg - tmpRoi_3_avg))
        sub24 = np.sum((tmpRoi_2_avg - tmpRoi_4_avg) * (tmpRoi_2_avg - tmpRoi_4_avg))

        # print(sub12, sub14, sub23, sub24)
        if sub14 > subThresold[0] or sub23 > subThresold[1] or sub12 < subThresold[2] or sub24 < subThresold[3]:
            return False
        if SHOW_SAVE_IMG:
            dirPath = os.getcwd() + "/rois/reg"
            if not os.path.exists(dirPath):
                os.mkdir(dirPath)
            cv2.imwrite(dirPath + "/" + "roi%d - %d.jpg" % (center[0], center[1]), binaryImg)
        return True
    return False

def PreDeal(imgPath):
    """
    图像预处理；
    :param path:
    :return:
    """
    img = cv2.imread(imgPath)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray = np.float32(gray)

    # 输入图像必须是 float32 ,最后一个参数在 0.04 到 0.05 之间
    dst = cv2.cornerHarris(gray, 4, 3, 0.05)

    # result is dilated for marking the corners, not important
    dst = cv2.dilate(dst, None)

    tmpNewImg = np.zeros((img.shape[0], img.shape[1], 1), dtype="uint8")

    # Threshold for an optimal value, it may vary depending on the image.
    tmpNewImg[dst > ratio * dst.max()] = 255

    if SHOW_SAVE_IMG:
        dirPath = os.getcwd() + "/rois/reg"
        if os.path.exists(dirPath):
            shutil.rmtree(dirPath)
        # os.mkdir(dirPath)

    # 平滑；
    img = cv2.medianBlur(img, 3)

    newContours = []
    contours, hierarchy = cv2.findContours(tmpNewImg.copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    nContours = len(contours)
    regs = []
    for i in range(nContours):
        if contours[i].size < 100:                          # 50个点一下的进入识别；
            newContours.append(contours[i])
            bReg = AnalysisRegNear(contours[i], img)           # 根据灰度、颜色特征识别 TOOD: 还可以进行外部轮廓辅助识别；
            if bReg:
                regs.append(contours[i])                # reg位置；
                if SHOW_SAVE_IMG:
                    cv2.drawContours(img, contours[i], -1, (0, 0, 255), 10)

    # 显示定位到的reg标记；
    if SHOW_SAVE_IMG:
        cv2.namedWindow("enhanced", 0)
        cv2.resizeWindow("enhanced", 640, 480)
        cv2.imshow("enhanced", img)
        cv2.waitKey(0)

        # 显示所有的符合标准的轮廓；
        cv2.drawContours(img, newContours, -1, (0, 0, 255), 3)
        cv2.namedWindow("filter", 0)
        cv2.resizeWindow("filter", 640, 480)
        cv2.imshow("filter", img)
        cv2.waitKey(0)
    return regs                     # 返回所有识别的reg， 有可能失败，  但是目前还没有办法进行纠偏；

if __name__ == "__main__":
    tmpPath = os.getcwd() + "/IMG_20210315_072108.jpg"
    PreDeal(tmpPath)